import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
import warnings
warnings.filterwarnings('ignore')

# 导入数据预处理类
from sales_prediction_models import DataPreprocessor

class PredictionEnsembler:
    """预测结果整合类"""
    
    def __init__(self):
        self.predictions = {}
        self.weights = {}
        self.ensemble_method = 'weighted_average'
        
    def add_predictions(self, model_name, predictions, weight=None):
        """添加单个模型的预测结果"""
        self.predictions[model_name] = np.array(predictions)
        if weight is not None:
            self.weights[model_name] = weight
            
    def set_weights(self, weights_dict):
        """设置模型权重"""
        self.weights = weights_dict.copy()
        
    def ensemble_predict(self, method='weighted_average'):
        """
        整合多个模型的预测结果
        
        方法:
        - simple_average: 简单平均
        - weighted_average: 加权平均
        - median: 中位数
        - trimmed_mean: 截尾平均
        """
        if not self.predictions:
            raise ValueError("没有可用的预测结果")
            
        predictions_df = pd.DataFrame(self.predictions)
        
        if method == 'simple_average':
            return predictions_df.mean(axis=1).values
            
        elif method == 'weighted_average':
            if not self.weights:
                # 如果没有指定权重，使用简单平均
                return predictions_df.mean(axis=1).values
            
            # 确保所有权重的和为1
            total_weight = sum(self.weights.values())
            normalized_weights = {k: v/total_weight for k, v in self.weights.items()}
            
            # 加权平均
            weighted_pred = np.zeros(len(predictions_df))
            for model_name, weight in normalized_weights.items():
                if model_name in predictions_df.columns:
                    weighted_pred += predictions_df[model_name] * weight
            
            return weighted_pred
            
        elif method == 'median':
            return predictions_df.median(axis=1).values
            
        elif method == 'trimmed_mean':
            # 去除最高和最低各10%后取平均
            return predictions_df.apply(lambda x: self._trimmed_mean(x, 0.1), axis=1).values
            
        else:
            raise ValueError(f"不支持的整合方法: {method}")
    
    def _trimmed_mean(self, series, trim_ratio):
        """计算截尾平均"""
        n = len(series)
        trim_count = int(n * trim_ratio)
        sorted_values = np.sort(series)[trim_count:n-trim_count]
        return np.mean(sorted_values) if len(sorted_values) > 0 else np.mean(series)
    
    def get_prediction_stats(self):
        """获取预测结果的统计信息"""
        if not self.predictions:
            return {}
            
        pred_df = pd.DataFrame(self.predictions)
        stats = {
            'mean': pred_df.mean(axis=1).values,
            'std': pred_df.std(axis=1).values,
            'min': pred_df.min(axis=1).values,
            'max': pred_df.max(axis=1).values,
            'median': pred_df.median(axis=1).values
        }
        
        # 添加变异系数
        stats['cv'] = stats['std'] / (stats['mean'] + 1e-8)  # 避免除零
        
        return stats

class DataPostProcessor:
    """数据后处理类"""
    
    def __init__(self):
        self.processing_rules = {
            'min_value': 0,      # 最小值约束
            'max_value': None,   # 最大值约束
            'integer': True,     # 是否转换为整数
            'smooth': False,     # 是否平滑处理
            'outlier_handle': True  # 是否处理异常值
        }
    
    def set_processing_rules(self, rules_dict):
        """设置处理规则"""
        self.processing_rules.update(rules_dict)
    
    def post_process(self, predictions, historical_data=None):
        """
        后处理预测结果
        
        Args:
            predictions: 原始预测值
            historical_data: 历史数据，用于计算统计约束
        """
        processed = np.array(predictions).copy()
        
        # 1. 处理负值
        if self.processing_rules['min_value'] is not None:
            processed = np.maximum(processed, self.processing_rules['min_value'])
        
        # 2. 处理异常值（基于历史数据的IQR方法）
        if self.processing_rules['outlier_handle'] and historical_data is not None:
            q1, q3 = np.percentile(historical_data, [25, 75])
            iqr = q3 - q1
            lower_bound = q1 - 1.5 * iqr
            upper_bound = q3 + 1.5 * iqr
            
            # 将异常值限制在合理范围内
            processed = np.clip(processed, lower_bound, upper_bound)
        
        # 3. 最大值约束
        if self.processing_rules['max_value'] is not None:
            processed = np.minimum(processed, self.processing_rules['max_value'])
        
        # 4. 整数化处理（适用于销售数量预测）
        if self.processing_rules['integer']:
            processed = np.round(processed).astype(int)
        
        # 5. 平滑处理（移动平均）
        if self.processing_rules['smooth'] and len(processed) > 3:
            window_size = min(3, len(processed) // 2)
            processed = self._moving_average(processed, window_size)
        
        return processed
    
    def _moving_average(self, data, window_size):
        """移动平均平滑"""
        if len(data) <= window_size:
            return data
            
        smoothed = data.copy()
        half_window = window_size // 2
        
        for i in range(half_window, len(data) - half_window):
            smoothed[i] = np.mean(data[i-half_window:i+half_window+1])
        
        return smoothed
    
    def apply_business_rules(self, predictions, dates=None):
        """
        应用业务规则
        
        例如：周末销售通常较低，节假日销售较高等
        """
        processed = np.array(predictions).copy()
        
        if dates is not None:
            dates = pd.to_datetime(dates)
            
            # 周末销售通常较低（假设为工作日的70%）
            is_weekend = dates.dt.dayofweek.isin([5, 6])  # 周六、周日
            processed[is_weekend] = processed[is_weekend] * 0.7
            
            # 月初和月末可能销售较高
            is_month_start = dates.dt.day.isin([1, 2, 3])
            is_month_end = dates.dt.day.isin([28, 29, 30, 31])
            processed[is_month_start | is_month_end] = processed[is_month_start | is_month_end] * 1.2
        
        return processed

class AccuracyEvaluator:
    """准确率评估类"""
    
    def __init__(self):
        self.metrics = {}
        
    def calculate_metrics(self, y_true, y_pred, zero_threshold=1e-6):
        """
        计算各种评估指标
        
        Args:
            y_true: 真实值
            y_pred: 预测值
            zero_threshold: 判断为0的阈值
        """
        y_true = np.array(y_true)
        y_pred = np.array(y_pred)
        
        # 基础指标
        mae = mean_absolute_error(y_true, y_pred)
        mse = mean_squared_error(y_true, y_pred)
        rmse = np.sqrt(mse)
        
        # 处理可能为0的情况
        mask = y_true > zero_threshold
        
        if np.sum(mask) > 0:
            mape = np.mean(np.abs((y_true[mask] - y_pred[mask]) / y_true[mask])) * 100
        else:
            mape = np.inf
        
        # R²分数
        try:
            r2 = r2_score(y_true, y_pred)
        except:
            r2 = np.nan
        
        # 自定义指标：零值预测准确率
        true_zeros = y_true <= zero_threshold
        pred_zeros = y_pred <= zero_threshold
        zero_accuracy = np.mean(true_zeros == pred_zeros) * 100
        
        # 非零值预测准确率
        non_zero_mask = y_true > zero_threshold
        if np.sum(non_zero_mask) > 0:
            non_zero_mae = mean_absolute_error(y_true[non_zero_mask], y_pred[non_zero_mask])
            non_zero_mape = np.mean(np.abs((y_true[non_zero_mask] - y_pred[non_zero_mask]) / 
                                         y_true[non_zero_mask])) * 100
        else:
            non_zero_mae = np.nan
            non_zero_mape = np.nan
        
        metrics = {
            'MAE': mae,
            'MSE': mse,
            'RMSE': rmse,
            'MAPE': mape,
            'R2': r2,
            'Zero_Accuracy': zero_accuracy,
            'Non_Zero_MAE': non_zero_mae,
            'Non_Zero_MAPE': non_zero_mape
        }
        
        self.metrics = metrics
        return metrics
    
    def print_evaluation_report(self):
        """打印评估报告"""
        if not self.metrics:
            print("没有可用的评估指标")
            return
            
        print("=== 模型评估报告 ===")
        print(f"平均绝对误差 (MAE): {self.metrics['MAE']:.4f}")
        print(f"均方误差 (MSE): {self.metrics['MSE']:.4f}")
        print(f"均方根误差 (RMSE): {self.metrics['RMSE']:.4f}")
        print(f"平均绝对百分比误差 (MAPE): {self.metrics['MAPE']:.2f}%")
        print(f"R² 分数: {self.metrics['R2']:.4f}")
        print(f"零值预测准确率: {self.metrics['Zero_Accuracy']:.2f}%")
        
        if not np.isnan(self.metrics['Non_Zero_MAE']):
            print(f"非零值平均绝对误差: {self.metrics['Non_Zero_MAE']:.4f}")
            print(f"非零值平均绝对百分比误差: {self.metrics['Non_Zero_MAPE']:.2f}%")
    
    def compare_models(self, results_dict):
        """
        比较多个模型的结果
        
        Args:
            results_dict: 字典，键为模型名称，值为(y_true, y_pred)元组
        """
        comparison_df = pd.DataFrame()
        
        for model_name, (y_true, y_pred) in results_dict.items():
            metrics = self.calculate_metrics(y_true, y_pred)
            comparison_df[model_name] = pd.Series(metrics)
        
        return comparison_df

class SalesPredictionPipeline:
    """销售预测完整流程类"""
    
    def __init__(self):
        self.models = {}
        self.preprocessor = DataPreprocessor()
        self.ensembler = PredictionEnsembler()
        self.post_processor = DataPostProcessor()
        self.evaluator = AccuracyEvaluator()
        self.historical_data = None
        
    def add_model(self, model_name, model_instance):
        """添加模型"""
        self.models[model_name] = model_instance
        
    def train_all_models(self, X_train, y_train):
        """训练所有模型"""
        training_results = {}
        
        for model_name, model in self.models.items():
            print(f"正在训练 {model_name}...")
            try:
                model.train(X_train, y_train)
                training_results[model_name] = "成功"
            except Exception as e:
                training_results[model_name] = f"失败: {str(e)}"
                print(f"训练 {model_name} 失败: {e}")
        
        return training_results
    
    def predict_with_all_models(self, X_test):
        """使用所有模型进行预测"""
        all_predictions = {}
        
        for model_name, model in self.models.items():
            if model.is_trained:
                try:
                    predictions = model.predict(X_test)
                    all_predictions[model_name] = predictions
                except Exception as e:
                    print(f"模型 {model_name} 预测失败: {e}")
        
        return all_predictions
    
    def run_ensemble_prediction(self, all_predictions, method='weighted_average'):
        """运行集成预测"""
        # 清空之前的预测结果
        self.ensembler.predictions.clear()
        self.ensembler.weights.clear()
        
        # 添加所有预测结果
        for model_name, predictions in all_predictions.items():
            self.ensembler.add_predictions(model_name, predictions)
        
        # 设置权重（这里可以根据模型在历史数据上的表现来设置）
        # 简单起见，使用等权重
        if not self.ensembler.weights:
            n_models = len(all_predictions)
            for model_name in all_predictions.keys():
                self.ensembler.weights[model_name] = 1.0 / n_models
        
        # 执行集成预测
        ensemble_pred = self.ensembler.ensemble_predict(method)
        
        return ensemble_pred